package com.assistant.audio;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.ArrayList;
import java.util.List;


/**
 * @author Anyang
 * @version 1.0.0
 * @date 2025/10/25 05:10:41
 * @description VAD 模块：封装了基于能量的语音活动检测逻辑、音频缓冲及格式转换。
 */
public class VADModule {
    private static final Logger logger = LoggerFactory.getLogger(VADModule.class);
    private final double energyThreshold;
    private final int silenceFramesRequired;
    
    private int framesOfSilence;
    private final List<Short> audioBuffer;

    /**
     * @param energyThreshold 能量阈值，低于此值为静音
     * @param silenceFramesRequired 判定为静音所需的连续帧数
     */
    public VADModule(double energyThreshold, int silenceFramesRequired) {
        this.energyThreshold = energyThreshold;
        this.silenceFramesRequired = silenceFramesRequired;
        this.audioBuffer = new ArrayList<>();
        this.framesOfSilence = 0;
        logger.info("VAD 模块已初始化：能量阈值={}, 静音帧数={}", energyThreshold, silenceFramesRequired);
    }

    /**
     * 重置 VAD 状态，清空缓冲区和静音计数。
     * (在每次唤醒时调用)
     */
    public void reset() {
        this.framesOfSilence = 0;
        this.audioBuffer.clear();
        logger.debug("VAD 状态已重置。");
    }

    /**
     * 处理传入的音频帧。
     * @param frame 512 short (来自麦克风)
     * @return true 如果检测到语音结束 (达到静音阈值)，否则返回 false。
     */
    public boolean processFrame(short[] frame) {
        // 1. 缓冲音频
        // (注意：我们将所有音频都缓冲，包括静音，STT 对末尾静音有一定鲁棒性)
        for (short s : frame) {
            this.audioBuffer.add(s);
        }

        // 2. 计算能量
        double rms = calculateRMS(frame);
        // logger.trace("VAD Energy (RMS): {}", rms); // 调试时使用

        // 3. VAD 逻辑
        if (rms > this.energyThreshold) {
            framesOfSilence = 0; // 正在说话
        } else {
            framesOfSilence++; // 可能是静音
        }

        // 4. 判断是否结束
        if (framesOfSilence >= this.silenceFramesRequired) {
            logger.info("【VAD】检测到指令结束 (静音 {} 帧)。", framesOfSilence);
            return true;
        }
        
        return false;
    }

    /**
     * 获取缓冲的音频数据 (PCM, 16kHz, 16bit, mono)，并清空缓冲区。
     * @return 包含所有录音的 byte[]
     */
    public byte[] getAndClearAudioBuffer() {
        if (this.audioBuffer.isEmpty()) {
            return new byte[0];
        }
        
        // 复制并转换
        byte[] audioData = convertShortsToBytes(this.audioBuffer);
        logger.info("VAD 已提取音频数据 {} 字节 ({} 采样点)", audioData.length, this.audioBuffer.size());
        
        // 清空
        this.audioBuffer.clear();
        this.framesOfSilence = 0; // 顺便重置
        
        return audioData;
    }

    // --- 私有辅助工具方法 (从 MainApplication 迁移而来) ---

    /**
     * 计算音频帧的 RMS 能量（音量）
     */
    private double calculateRMS(short[] frame) {
        long sumOfSquares = 0;
        for (short s : frame) {
            sumOfSquares += (long) s * s;
        }
        return Math.sqrt((double) sumOfSquares / frame.length);
    }

    private byte[] convertShortsToBytes(List<Short> shorts) {
        byte[] bytes = new byte[shorts.size() * 2];
        ByteBuffer.wrap(bytes)
                .order(ByteOrder.LITTLE_ENDIAN)
                .asShortBuffer()
                .put(toPrimitiveArray(shorts));
        return bytes;
    }

    private short[] toPrimitiveArray(List<Short> shorts) {
        short[] arr = new short[shorts.size()];
        for (int i = 0; i < shorts.size(); i++) {
            arr[i] = shorts.get(i);
        }
        return arr;
    }
    /**
     * 检查 VAD 是否处于“正在录音”状态（即 audioBuffer 不为空）。
     * 这是判断用户是否在说话的关键逻辑。
     * @return 如果缓冲区中有数据则返回 true。
     */
    public boolean isSpeaking() {
        // 只要音频缓冲区不为空，就认为 VAD 已经捕捉到语音，处于“正在说话/录音”状态。
        return !this.audioBuffer.isEmpty();
    }
}